import random

import numpy as np
import torch
from torchvision import transforms as T
from torchvision.transforms import functional as F


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class RandomCrop:
    def __init__(self, w, h):
        self.w = w
        self.h = h

    def __call__(self, image, target):
        crop_params = T.RandomCrop.get_params(image, (self.h, self.w))
        image = F.crop(image, *crop_params)
        if target is not None:
            target = F.crop(target, *crop_params)
        return image, target


class ToPILImage:
    def __init__(self):
        self.trans = T.ToPILImage()

    def __call__(self, image, target):
        image = self.trans(image)
        target = self.trans(target)
        return image, target


class ToTensor:
    def __init__(self):
        self.trans = T.ToTensor()

    def __call__(self, image, target):
        image = self.trans(image)
        target = self.trans(target)
        return image, target


class RandomRotation:
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, image, target):
        image = F.rotate(image, self.degree)
        target = F.rotate(target, self.degree)
        return image, target


class RandomHorizontalFlip:
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            target = F.hflip(target)
        return image, target


class RandomVerticalFlip:
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.vflip(image)
            target = F.vflip(target)
        return image, target


class RandomResizedCrop:
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        crop_params = T.RandomResizedCrop.get_params(image, scale=[0.08, 1.0], ratio=[0.75, 1.3333333333333333])
        image = F.crop(image, *crop_params)
        target = F.crop(target, *crop_params)
        image = F.resize(image, [self.size, self.size])
        target = F.resize(target, [self.size, self.size])
        return image, target


class GaussianBlur:
    def __init__(self, kernel_size):
        self.kernel_size = kernel_size

    def __call__(self, image, target):
        image = F.gaussian_blur(image, self.kernel_size)
        return image, target


class ColorJit:
    def __call__(self, image, target):
        jitter = T.ColorJitter.get_params((0.8, 1.2), (0.8, 1.2), (0.8, 1.2), (-0.1, 0.1))
        image = jitter(image)
        return image, target


class Normalize:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, target



